from __future__ import division
import numpy as np
from scipy.sparse import coo_matrix
from scipy.sparse.linalg import spsolve
import nlopt
import matplotlib
from matplotlib import colors
import matplotlib.pyplot as plt

import time
#from pyomo.environ import *
#from pyomo.opt import SolverFactory


# 输入参数
class para:
    # xyz方向上单元格数量
    nelx = 30
    nely = 10
    nelz = 5
    # 体积分数
    volfrac = 0.5
    # 滤波半径
    rmin = 1.2
    # SIMP 惩罚因子
    penal = 3.0
    # 滤波类型
    ft = 0

class eMat:
    def __init__(self,nelx,nely,nelz):
        self.nelx=nelx
        self.nely=nely
        self.nelz=nelz
    def Mat(self):
        nelx=self.nelx
        nely=self.nely
        nelz=self.nelz
        Mat = np.zeros((nelx * nely * nelz, 24), dtype=int)
        for elz in range(nelz):
            for elx in range(nelx):
                for ely in range(nely):
                    el = ely + (elx * nely) + elz * (nelx * nely)
                    n1 = elz * (nelx + 1) * (nely + 1) + (nely + 1) * elx + ely
                    n2 = elz * (nelx + 1) * (nely + 1) + (nely + 1) * (elx + 1) + ely
                    n3 = (elz + 1) * (nelx + 1) * (nely + 1) + (nely + 1) * elx + ely
                    n4 = (elz + 1) * (nelx + 1) * (nely + 1) + (nely + 1) * (elx + 1) + ely
                    Mat[el, :] = np.array(
                        [3 * n1 + 3, 3 * n1 + 4, 3 * n1 + 5, 3 * n2 + 3, 3 * n2 + 4, 3 * n2 + 5, \
                         3 * n2, 3 * n2 + 1, 3 * n2 + 2, 3 * n1, 3 * n1 + 1, 3 * n1 + 2, \
                         3 * n3 + 3, 3 * n3 + 4, 3 * n3 + 5, 3 * n4 + 3, 3 * n4 + 4, 3 * n4 + 5, \
                         3 * n4, 3 * n4 + 1, 3 * n4 + 2, 3 * n3, 3 * n3 + 1, 3 * n3 + 2])
        return Mat

class FE:
    '''
    有限元分析，求解位移向量
    '''
    def __init__(self,x,nelx,nely,nelz,volfrac,rmin,penal,ft,Emin,Emax,KE):
        self.x=x
        self.nelx=nelx
        self.nely=nely
        self.nelz=nelz
        self.volfrac=volfrac
        self.rmin=rmin
        self.penal=penal
        self.ft=ft
        self.Emin=Emin
        self.Emax=Emax
        self.KE=KE

    '''
    在材料分布x下的结构位移向量u
    '''
    def Usolution(self):
        x=self.x
        nelx=self.nelx
        nely=self.nely
        nelz=self.nelz
        volfrac=self.volfrac
        rmin=self.rmin
        penal=self.penal
        ft=self.ft
        Emin=self.Emin
        Emax=self.Emax
        KE=self.KE
        # 自由度数目：每个节点3个自由度
        ndof = 3 * (nelx + 1) * (nely + 1) * (nelz + 1)
        # FE: Build the index vectors for the for coo matrix format.
        ee = eMat(nelx,nely,nelz)
        edofMat = ee.Mat()
        # Construct the index pointers for the coo format
        iK = np.kron(edofMat, np.ones((24, 1))).flatten()
        jK = np.kron(edofMat, np.ones((1, 24))).flatten()
        # USER - DEFINED LOAD DOFs
        kl = np.arange(nelz + 1)
        loadnid = kl * (nelx + 1) * (nely + 1) + (nely + 1) * (nelx + 1) - 1  # Node IDs
        loaddof = 3 * loadnid + 1  # DOFs
        # USER - DEFINED SUPPORT FIXED DOFs
        [jf, kf] = np.meshgrid(np.arange(nely + 1), np.arange(nelz + 1))  # Coordinates
        fixednid = (kf) * (nely + 1) * (nelx + 1) + jf  # Node IDs
        fixeddof = np.array([3 * fixednid, 3 * fixednid + 1, 3 * fixednid + 2]).flatten()  # DOFs
        # BC's and support
        dofs = np.arange(3 * (nelx + 1) * (nely + 1) * (nelz + 1))
        free = np.setdiff1d(dofs, fixeddof)
        # Solution and RHS vectors
        f = np.zeros((ndof, 1))
        u = np.zeros((ndof, 1))
        # Set load
        f[loaddof, 0] = -1
        # Setup and solve FE problem
        sK = ((KE.flatten()[np.newaxis]).T * (Emin + x ** penal * (Emax - Emin))).flatten(order='F')
        K = coo_matrix((sK, (iK, jK)), shape=(ndof, ndof)).tocsc()
        # Remove constrained dofs from matrix
        K = K[free, :][:, free]
        # 求解有限元系统
        u[free, 0] = spsolve(K, f[free, 0])
        return u

# Filter: Build (and assemble) the index+data vectors for the coo matrix format
class FILTERMATRIX:
    def __init__(self,nelx,nely,nelz,rmin):
        self.nelx=nelx
        self.nely=nely
        self.nelz=nelz
        self.rmin=rmin


    def assembly(self):
        nelx=self.nelx
        nely=self.nely
        nelz=self.nelz
        rmin=self.rmin
        nfilter = int(nelx * nely * nelz * ((2 * (np.ceil(rmin) - 1) + 1) ** 3))
        iH = np.zeros(nfilter)
        jH = np.zeros(nfilter)
        sH = np.zeros(nfilter)
        cc = 0
        for z in range(nelz):
            for i in range(nelx):
                for j in range(nely):
                    row = i * nely + j + z * (nelx * nely)
                    kk1 = int(np.maximum(i - (np.ceil(rmin) - 1), 0))
                    kk2 = int(np.minimum(i + np.ceil(rmin), nelx))
                    ll1 = int(np.maximum(j - (np.ceil(rmin) - 1), 0))
                    ll2 = int(np.minimum(j + np.ceil(rmin), nely))
                    mm1 = int(np.maximum(z - (np.ceil(rmin) - 1), 0))
                    mm2 = int(np.minimum(z + np.ceil(rmin), nelz))
                    for m in range(mm1, mm2):
                        for k in range(kk1, kk2):
                            for l in range(ll1, ll2):
                                col = k * nely + l + m * (nelx * nely)
                                fac = rmin - np.sqrt((i - k) * (i - k) + (j - l) * (j - l) + (z - m) * (z - m))
                                iH[cc] = row
                                jH[cc] = col
                                sH[cc] = np.maximum(0.0, fac)
                                cc = cc + 1
        # Finalize assembly and convert to csc format
        H = coo_matrix((sH, (iH, jH)), shape=(nelx * nely * nelz, nelx * nely * nelz)).tocsc()
        return H


#Optimization Problem Definition
def objfunc(x,grad):
    if grad.size > 0:
        # Default input parameters
        p = para()
        nelx = p.nelx
        nely = p.nely
        nelz = p.nelz
        volfrac = p.volfrac
        rmin = p.rmin
        penal = p.penal
        ft = p.ft  # ft==0 -> sens, ft==1 -> dens
        # Max and min stiffness
        Emin = 1e-3
        Emax = 1.0
        # list to array
        # Finalize assembly and convert to csc format
        HH = FILTERMATRIX(nelx, nely, nelz, rmin)
        H = HH.assembly()
        Hs = H.sum(1)
        # 刚体矩阵
        KE = lk()
        # FE
        uu = FE(x, nelx, nely, nelz, volfrac, rmin, penal, ft, Emin, Emax, KE)
        u = uu.Usolution()
        # sensitivity
        dv = np.ones(nely * nelx * nelz)
        dc = np.ones(nely * nelx * nelz)
        ce = np.ones(nely * nelx * nelz)
        ee = eMat(nelx, nely, nelz)
        edofMat = ee.Mat()
        ce[:] = (
        np.dot(u[edofMat].reshape(nelx * nely * nelz, 24), KE) * u[edofMat].reshape(nelx * nely * nelz, 24)).sum(1)
        dc[:] = (-penal * x ** (penal - 1) * (Emax - Emin)) * ce
        dv[:] = np.ones(nely * nelx * nelz)
        # Sensitivity filtering:
        if ft == 0:
            dc[:] = np.asarray((H * (x * dc))[np.newaxis].T / Hs)[:, 0] / np.maximum(0.001, x)
        elif ft == 1:
            dc[:] = np.asarray(H * (dc[np.newaxis].T / Hs))[:, 0]
            dv[:] = np.asarray(H * (dv[np.newaxis].T / Hs))[:, 0]
        # gradient of obj
        grad[:] = dc
    f = ((Emin + x ** penal * (Emax - Emin)) * ce).sum()
    return f

def Constraint1(x, grad):
    if grad.size > 0:
        p = para()
        nelx = p.nelx
        nely = p.nely
        nelz = p.nelz
        # gradient of con
        dg = list(1 / (0.5 * nelx * nely * nelz + 1e-2) for i in range(nelx * nely * nelz))
        grad[:] = np.array(dg)
    return x.sum()/(0.5 * nelx * nely * nelz + 1e-2)-1

def Constraint2(x, grad):
    if grad.size > 0:
        p = para()
        nelx = p.nelx
        nely = p.nely
        nelz = p.nelz
        # gradient of con
        dg = list(-1 / (0.5 * nelx * nely * nelz - 1e-2) for i in range(nelx * nely * nelz))
        grad[:] = np.array(dg)
    return -x.sum()/(0.5 * nelx * nely * nelz - 1e-2)+1

# 主函数
def main():
    # Default input parameters
    p=para()
    nelx = p.nelx
    nely = p.nely
    nelz = p.nelz
    volfrac = p.volfrac
    rmin = p.rmin
    penal = p.penal
    ft = p.ft # ft==0 -> sens, ft==1 -> dens
    # 材料刚度的最大值和最小值
    Emin = 1e-3
    Emax = 1.0
    # 迭代计数器
    loop = 0
    # 迭代变化量
    change = 1
    # 密度数组：三维展成一维，每个单元的材料密度，初始为体积密度
    xx = volfrac * np.ones(nely * nelx * nelz, dtype=float)

    # 启用 matplotlib 的交互模式，使得后续图像可以实时更新，不阻塞程序执行
    # plt.ion()

    # 列表存储对象
    figures = []
    axes = []
    images = []

    # 初始结构可视化：绘制沿 z 方向的横截面（即 xy 平面上每一层）
    for i in range(nelz):
        # locals()：返回当前局部作用域中的变量字典，{局部变量名：值}
        # fig：整个画布
        # ax：画布上子图
        # fig：子图显示的具体像素矩阵
        # locals()['fig' + str(i+1)], locals()['ax'+str(i+1)] = plt.subplots()
        # # z方向上每一层的密度值：xx三维展平成一维的密度数组
        # # 一维数组 reshape成（x,y）
        # locals()['im' + str(i+1)] = locals()['ax'+str(i+1)].imshow(-xx[nelx*nely*i:nelx*nely*(i+1)].reshape((nelx, nely)).T, cmap='gray', \
        #              interpolation='none', norm=colors.Normalize(vmin=-1, vmax=0))
        # locals()['fig' + str(i + 1)].show()
        # plt.show()

        fig,ax = plt.subplots(figsize=(6,4))
        im = ax.imshow(-xx[nelx*nely*i:nelx*nely*(i+1)].reshape((nelx, nely)).T, cmap='gray',interpolation='none',norm=colors.Normalize(vmin=-1, vmax=0))

        ax.set_title(f'Layer {i+1}')
        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        fig.colorbar(im, ax=ax)

        # 将图像对象添加到列表
        figures.append(fig)
        axes.append(ax)
        images.append(im)

        # 显示图像
        fig.show()

    t1 = 0

    # 优化终止条件
    while change > 0.001 and loop < 200:

        # told：上一次迭代开始时间，用于评估单次迭代的计算耗时
        t1old=t1
        #t1 = time.clock()
        # t1：当前迭代开始时间
        t1 = time.perf_counter()
        print(f"t1:-------------{t1}")

        # 上一次迭代旧解
        xxold = xx.copy()
        # 更新迭代次数
        loop = loop + 1

        # 调用优化器
        opt = nlopt.opt(nlopt.LD_MMA, nelx * nely * nelz)
#        opt = SolverFactory('ipopt')
#        opt = opt.solve(LD_MMA, nelx * nely * nelz)
#        opt = opt.solve(LD_MMA, nelx * nely * nelz)
        # 最小化目标函数
        opt.set_min_objective(objfunc)
        # 变量边界
        opt.set_lower_bounds(Emin)
        opt.set_upper_bounds(Emax)
        # 非线性约束
        opt.add_inequality_constraint(Constraint1, 1e-4)
        opt.add_inequality_constraint(Constraint2, 1e-4)
        # 停止条件
        opt.set_ftol_rel(1e-4)

        # Performing the optimization
        x = opt.optimize(xx)

        #t2=time.clock()
        # t2：调用优化器 opt.optimize(xx) 之后记录的时间点，t2 - t1 表示优化器执行优化计算所花费的时间
        t2=time.perf_counter()

        # Filter design variables
        # Finalize assembly and convert to csc format
        HH = FILTERMATRIX(nelx, nely, nelz, rmin)
        H = HH.assembly()
        Hs = H.sum(1)
        if ft == 0:
            xx = x
        elif ft == 1:
            xx = np.asarray(H * x[np.newaxis].T / Hs)[:, 0]
        # 当前最优目标函数值
        obj = opt.last_optimum_value()
        # vol: 当前材料体积分数
        vol = xx.sum() / (nelx * nely * nelz)
        # 变量变化量：最大绝对值
        change = np.linalg.norm(xx.reshape(nelx * nely * nelz, 1) - xxold.reshape(nelx * nely * nelz, 1), np.inf)
        # t3：在完成滤波设计变量之后记录的时间点，t3 - t2 表示包括滤波所花费的时间
        t3=time.perf_counter()
        #t3=time.clock()

        # Plot to screen
        for i in range(nelz):
            # locals()['im' + str(i + 1)].set_array(-x[nelx*nely*i:nelx*nely*(i+1)].reshape((nelx, nely)).T)
            # locals()['fig' + str(i + 1)].canvas.draw()
            images[i].set_array(-xx[nelx*nely*i:nelx*nely*(i+1)].reshape((nelx, nely)).T)
            figures[i].canvas.draw()
            figures[i].canvas.flush_events()  # 强制刷新图像
            plt.pause(0.001)

        #t4=time.clock()
        # t4：在更新图像数据并调用 canvas.draw() 等操作之后记录的时间点，t4 - t3 表示图像更新显示相关操作所花费的时间。
        t4=time.perf_counter()

        # Write iteration history to screen (req. Python 2.6 or newer)
        print("it.: {0} , obj.: {1:.3f} Vol.: {2:.3f}, ch.: {3:.3f}".format( \
            loop, obj, vol, change))
        print("t1-t1old.: {0:.5f} , t2-t1.: {1:.5f} , t3-t2.: {2:.5f} , t3-t2.: {3:.5f}". \
            format(t1-t1old, t2-t1, t3-t2, t4-t3))
    # Make sure the plot stays and that the shell remains
    plt.ioff()
    plt.show()

    #input("Press any key...")

# 返回一个立方体单元的 24x24 刚度矩阵
def lk():
    '''
    刚度矩阵：
    单元在形变下的力响应：三维问题，每个单元有8个节点，每个节点有3个自由度，刚度矩阵是24*24
    :return:
    '''
    # 泊松比
    nu = 0.3
    # 生成刚度参数k
    A = np.array([[32,6,-8,6,-6,4,3,-6,-10,3,-3,-3,-4,-8],[-48,0,0,-24,24,0,0,0,12,-12,0,12,12,12]])
    b = np.array([[1],[nu]])
    k = 1/float(144)*np.dot(A.T,b).flatten()

    K1 = np.array([[k[0],k[1],k[1],k[2],k[4],k[4]],
    [k[1],k[0],k[1],k[3],k[5],k[6]],
    [k[1],k[1],k[0],k[3],k[6],k[5]],
    [k[2],k[3],k[3],k[0],k[7],k[7]],
    [k[4],k[5],k[6],k[7],k[0],k[1]],
    [k[4],k[6],k[5],k[7],k[1],k[0]]])

    K2 = np.array([[k[8],k[7],k[11],k[5],k[3],k[6]],
    [k[7],k[8],k[11],k[4],k[2],k[4]],
    [k[9],k[9],k[12],k[6],k[3],k[5]],
    [k[5],k[4],k[10],k[8],k[1],k[9]],
    [k[3],k[2],k[4],k[1],k[8],k[11]],
    [k[10],k[3],k[5],k[11],k[9],k[12]]])

    K3 = np.array([[k[5],k[6],k[3],k[8],k[11],k[7]],
    [k[6],k[5],k[3],k[9],k[12],k[9]],
    [k[4],k[4],k[2],k[7],k[11],k[8]],
    [k[8],k[9],k[1],k[5],k[10],k[4]],
    [k[11],k[12],k[9],k[10],k[5],k[3]],
    [k[1],k[11],k[8],k[3],k[4],k[2]]])

    K4 = np.array([[k[13],k[10],k[10],k[12],k[9],k[9]],
    [k[10],k[13],k[10],k[11],k[8],k[7]],
    [k[10],k[10],k[13],k[11],k[7],k[8]],
    [k[12],k[11],k[11],k[13],k[6],k[6]],
    [k[9],k[8],k[7],k[6],k[13],k[10]],
    [k[9],k[7],k[8],k[6],k[10],k[13]]])

    K5 = np.array([[k[0],k[1],k[7],k[2],k[4],k[3]],
    [k[1],k[0],k[7],k[3],k[5],k[10]],
    [k[7],k[7],k[0],k[4],k[10],k[5]],
    [k[2],k[3],k[4],k[0],k[7],k[1]],
    [k[4],k[5],k[10],k[7],k[0],k[7]],
    [k[3],k[10],k[5],k[1],k[7],k[0]]])

    K6 = np.array([[k[13],k[10],k[6],k[12],k[9],k[11]],
    [k[10],k[13],k[6],k[11],k[8],k[1]],
    [k[6],k[6],k[13],k[9],k[1],k[8]],
    [k[12],k[11],k[9],k[13],k[6],k[10]],
    [k[9],k[8],k[1],k[6],k[13],k[6]],
    [k[11],k[1],k[8],k[10],k[6],k[13]]])

    KE1=np.hstack((K1,K2,K3,K4))
    KE2=np.hstack((K2.T,K5,K6,K3.T))
    KE3=np.hstack((K3.T,K6,K5.T,K2.T))
    KE4=np.hstack((K4,K3,K2,K1.T))
    KE = 1/float(((nu+1)*(1-2*nu)))*np.vstack((KE1,KE2,KE3,KE4))

    # 返回单元刚度矩阵：24*24 数组
    return(KE)

# The real main driver
if __name__ == "__main__":
    main()